import pywt
import numpy as np
import matplotlib.pyplot as plt

# 工厂方法生成不同类型信号
def signal_factory(signal_type, x):
    if signal_type == 'sin':
        return np.sin(20 * np.pi * x)
    elif signal_type == 'step':
        return np.where(x < 0.5, 1, -1)
    elif signal_type == 'spike':
        s = np.zeros_like(x)
        s[len(x)//2] = 1
        return s
    elif signal_type == 'chirp':
        return np.sin(50 * np.pi * x**2)
    else:
        raise ValueError('未知信号类型')

# 小波列表和信号类型
wavelets = ['haar', 'db2', 'sym4', 'coif1']
signal_types = ['sin', 'step', 'spike', 'chirp']

x = np.linspace(0, 1, num=1024)
fig, axes = plt.subplots(len(signal_types), len(wavelets)+1, figsize=(5*(len(wavelets)+1), 3*len(signal_types)))
for row, sig_type in enumerate(signal_types):
    signal = signal_factory(sig_type, x)
    axes[row, 0].plot(signal)
    axes[row, 0].set_title('原始信号', fontproperties="SimHei")
    if row == 0:
        axes[row, 0].set_ylabel(sig_type, fontproperties="SimHei")
    else:
        axes[row, 0].set_ylabel(sig_type, fontproperties="SimHei")
    for col, w in enumerate(wavelets):
        cA, cD = pywt.dwt(signal, w)
        axes[row, col+1].plot(cA, label='cA')
        axes[row, col+1].plot(cD, label='cD')
        axes[row, col+1].set_title(f"小波: {w}", fontproperties="SimHei")
        axes[row, col+1].legend()
plt.tight_layout()
plt.show()